from mcp.server.fastmcp import FastMCP
from mcp.types import (
    Completion,
    CompletionArgument,
    CompletionContext,
    PromptReference,
    ResourceTemplateReference,
)

mcp = FastMCP(name="Completions Server")


@mcp.prompt(description="代码风格指导")
def code_style_guide(language: str, style: str) -> str:
    return f"请为{language}语言提供{style}风格的代码指导"


@mcp.resource("github://{owner}/{repo}")
def github_repo_info(owner: str, repo: str) -> str:
    return f"Github 仓库 {owner}/{repo}的信息"


# 使用mcp.completion装饰器注册补全器
@mcp.completion()
async def handle_completion(
    ref: PromptReference | ResourceTemplateReference,
    argument: CompletionArgument,
    context: CompletionContext | None,
) -> Completion | None:
    if isinstance(ref, PromptReference):
        if ref.name == "code_style_guide":
            if argument.name == "language":
                languages = ["python", "javascript", "java", "c++", "c#", "rust"]
                filtered_languages = [
                    lang
                    for lang in languages
                    if lang.lower().startswith(argument.value.lower())
                ]
                return Completion(values=filtered_languages, hasMore=False)
            elif argument.name == "style":
                styles = [
                    "clean",
                    "functional",
                    "object-oriented",
                    "procedural",
                    "declarative",
                ]
                filtered_styles = [
                    style
                    for style in styles
                    if style.lower().startswith(argument.value.lower())
                ]
                return Completion(values=filtered_styles, hasMore=False)
            else:
                return None
    elif isinstance(ref, ResourceTemplateReference):
        if ref.uri == "github://{owner}/{repo}":
            if argument.name == "owner":
                owners = ["microsoft", "google", "facebook", "modelcontextprotocol"]
                filtered_owners = [
                    owner
                    for owner in owners
                    if owner.lower().startswith(argument.value.lower())
                ]
                return Completion(values=filtered_owners, hasMore=False)
            elif argument.name == "repo":
                # 仓库名补全，需考虑 owner 上下文
                if context and context.arguments and context.arguments.get("owner"):
                    owner = context.arguments.get("owner")
                    # 针对不同 owner 提供不同仓库名列表
                    if owner == "microsoft":
                        repos = ["vscode", "typescript", "dotnet", "azure-sdk"]
                    elif owner == "google":
                        repos = ["tensorflow", "kubernetes", "go", "chromium"]
                    elif owner == "modelcontextprotocol":
                        repos = ["python-sdk", "servers", "specification", "website"]
                    else:
                        repos = ["main", "master", "develop", "feature"]
                    # 根据输入值过滤匹配的仓库名
                    filtered_repos = [
                        repo
                        for repo in repos
                        if repo.lower().startswith(argument.value.lower())
                    ]
                    # 返回补全结果
                    return Completion(
                        values=filtered_repos,
                        hasMore=False,
                    )
            else:
                return None
    else:
        return None


if __name__ == "__main__":
    mcp.run(transport="stdio")
